package cn.longwingstech.intelligence.agent.config;

import cn.longwingstech.intelligence.agent.app.config.EmbeddedConfigProperties;
import lombok.Data;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;

@Configuration
@ConfigurationProperties(prefix = "spring.ai.vectorstore.pgvector")
@Data
public class AgentConfig {

    private String jdbcUrl;
    private String username;
    private String password;
    private String database;

    public DataSource pgDataSource() {
        return DataSourceBuilder.create()
                .url(jdbcUrl)
                .username(username)
                .password(password)
                .build();
    }
    @Value("#{'${spring.ai.vectorstore.pgvector.table-name}'}")
    private String vectorTableName;

    @Bean("openAiEmbeddingModel")
    public OpenAiEmbeddingModel openAiEmbeddingModel(EmbeddedConfigProperties embeddedConfigProperties) {
        OpenAiApi openAiApi = OpenAiApi.builder()
                .apiKey(embeddedConfigProperties.getApiKey())
                .baseUrl(embeddedConfigProperties.getBaseUrl())
                .build();
        OpenAiEmbeddingOptions embeddingOptions = OpenAiEmbeddingOptions
                .builder().model(embeddedConfigProperties.getModel()).build();
        return new OpenAiEmbeddingModel(openAiApi, MetadataMode.ALL, embeddingOptions);
    }
    @Bean
    public PgVectorStore pgVectorStore( OpenAiEmbeddingModel openAiEmbeddingModel) {
        JdbcTemplate jdbcTemplate = new JdbcTemplate(pgDataSource());
        return PgVectorStore.builder(jdbcTemplate,openAiEmbeddingModel)
                .vectorTableName(vectorTableName)
                .build();
    }

}
